// Date:   Thu Aug 14 04:21:23 PM 2025
// Mail:   lunar_ubuntu@qq.com
// Author: https://github.com/xiaoqixian

#pragma once

#include <atomic>
#include <cstddef>
#include <future>
#include <opencv2/videoio.hpp>
#include <string>
#include <thread>
#include <unordered_map>
#include <vector>

#include "context.hpp"
#include "frames.hpp"
#include "config.hpp"
#include "av.hpp"
#include "datapipe.hpp"

class Sampler;

struct VideoStream {
  Config config_;
  std::string url_;
  int label_ = -1;
  int src_idx_ = -1;

  AVFormatContext* fmt_ctx_ = nullptr; // Format context (demux)
  AVCodecParameters* codec_par_ = nullptr;
  AVCodec* codec_ = nullptr;      // Decoder
  std::unordered_map<int, AVCodecCtxWrapper> codec_ctx_map_;
  std::shared_mutex codec_ctx_mtx_;
  int video_stream_index_ = -1;   // Index of the video stream

  AVPixelFormat pix_fmt_;         // Pixel format of decoded frames

  SwsContext* sws_ctx_ = nullptr; // For pixel format conversion / resize

  VideoStream() noexcept = default;
  // Constructor / Destructor
  VideoStream(Config);
  ~VideoStream();

  VideoStream(VideoStream const&) = delete;
  VideoStream& operator=(const VideoStream&) = delete;
  VideoStream(VideoStream&&) noexcept;
  VideoStream& operator=(VideoStream&& other) noexcept;

  // seek to the start of the source
  void rewind();
};

struct DecodeRequest {
  enum class ReqType {
    KeyFrame,
    Frames
  };
  FramePkts const* pkts = nullptr;
  ReqType req_type;
  // TODO: inlined vector optimization
  std::promise<std::vector<Frame>> res_promise;

  DecodeRequest(FramePkts const* fp, ReqType rt)
    : pkts(fp), req_type(rt) 
  {}
};


class AVDecoder {
  friend class Sampler;
public:
  AVDecoder(int id, Config, std::atomic_bool* stopped, DataPipe<DecodeRequest>*);
  AVDecoder(AVDecoder&&) noexcept;
  AVDecoder& operator=(AVDecoder&&) noexcept;
  ~AVDecoder();
private:
  static void worker(AVDecoder*, int id);
  Frame decode_keyframe(FramePkts const&);
  std::vector<Frame> decode_frames(FramePkts const&);

  int decoder_id_ {-1};
  std::vector<AVCodecCtxWrapper> codec_ctx_list_;
  std::atomic_bool* stopped_ = nullptr;
  std::thread worker_th_;
public:
  DataPipe<DecodeRequest>* req_ch_ = nullptr;
private:

  AVFrame* av_frame_ = nullptr;
};

class ChannelWorker {
public:
  ChannelWorker(
    int id,
    Config ch_cfg,
    std::unordered_map<std::string, Config> const& srcs_cfg,
    DataPipe<FramePkts>& window,
    std::unordered_map<int, std::atomic_size_t>& total_emit_cnt
  );

  ~ChannelWorker();

  void start();

private:
  Config config_;
  int channel_id_;
  std::vector<VideoStream> srcs_;
  std::unordered_map<int, size_t> emit_seq_;
  DataPipe<FramePkts>& window_;
  std::unordered_map<int, std::atomic_size_t>& total_emit_cnt_;
};

class Sampler {
public:
  Sampler(Config);

  void start();
  void wait();
  void stop_decoders();
  void notify_all_consumers();

  int get_pkts(int ch_id, size_t num, std::vector<FramePkts>& buf, std::atomic_bool&);

  std::future<std::vector<Frame>> decode_keyframe(FramePkts const&);
  std::future<std::vector<Frame>> decode_frames(FramePkts const&);

  Config config_;
  std::unordered_map<std::string, Config> srcs_config_;
  std::vector<DataPipe<FramePkts>> windows_;
  std::vector<std::thread> ch_workers_;

  std::vector<AVDecoder> decoders_;
  std::atomic_bool decoders_stopped_ {false};
  std::unordered_map<int, std::atomic_size_t> label_emit_cnt_;
  DataPipe<DecodeRequest> req_ch_;
};
